import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, losses, optimizers


class CAE(keras.Model):
    def __init__(self):
        super().__init__()
    # def build(self, input_shape):
        # self.inputs = layers.Input(shape=(5, 5, 1), name="inputs")
        self.cov_encoder = layers.Conv2D(15, (3, 3), strides=[1, 1], kernel_initializer=keras.initializers.glorot_uniform(seed=None),
                      activation='relu',padding="same")
        # self.maxpool_encoder=keras.layers.MaxPool2D(pool_size=(3,3), strides=[1, 1], padding="same")
        self.cov_middle = layers.Conv2D(30, (3, 3), strides=[1, 1], kernel_initializer=keras.initializers.glorot_uniform(seed=None),
                    activation='relu',padding="same",name="encoder")
        # self.maxpool_decoder =tf.nn.max_pool2d(pool_size=(3,3), strides=[1, 1], padding="same")
        self.covT_decoder = layers.Conv2DTranspose(15, (3, 3), strides=[1, 1],
                               kernel_initializer=keras.initializers.glorot_uniform(seed=None),
                                activation='relu',padding="same")
        self.covT_out_decoder = layers.Conv2DTranspose(1, (3, 3), strides=[1, 1],
                               kernel_initializer=keras.initializers.glorot_uniform(seed=None),
                               padding="same",name="decoder")
    def call(self,X,training=None, mask=None):
        # 编码
        self.out = self.cov_encoder(X)
        # self.out = self.maxpool_encoder(self.out)
        self.out = self.cov_middle(self.out)
        # self.out = self.maxpool_decoder(self.out)
        self.out = self.covT_decoder(self.out)
        out = self.covT_out_decoder(self.out)
        return out

    def middle_out(self,X):
        self.out = self.cov_encoder(X)
        # self.out = self.maxpool_encoder(self.out)
        middle_out = self.cov_middle(self.out)
        return middle_out



